package semantics;

import ic.LanguageException;
import ic.ast.BinaryOp;
import ic.ast.ClassType;
import ic.ast.DeclClass;
import ic.ast.DeclField;
import ic.ast.DeclLibraryMethod;
import ic.ast.DeclMethod;
import ic.ast.DeclStaticMethod;
import ic.ast.DeclVirtualMethod;
import ic.ast.Expression;
import ic.ast.HasName;
import ic.ast.HasType;
import ic.ast.Length;
import ic.ast.Literal;
import ic.ast.LocalVariable;
import ic.ast.NewArray;
import ic.ast.NewInstance;
import ic.ast.Node;
import ic.ast.Parameter;
import ic.ast.PrimitiveType;
import ic.ast.PrimitiveType.DataType;
import ic.ast.Program;
import ic.ast.Ref;
import ic.ast.RefArrayElement;
import ic.ast.RefField;
import ic.ast.RefVariable;
import ic.ast.StaticCall;
import ic.ast.StmtAssignment;
import ic.ast.StmtBlock;
import ic.ast.StmtBreak;
import ic.ast.StmtCall;
import ic.ast.StmtContinue;
import ic.ast.StmtIf;
import ic.ast.StmtReturn;
import ic.ast.StmtWhile;
import ic.ast.This;
import ic.ast.Type;
import ic.ast.UnaryOp;
import ic.ast.VirtualCall;
import ic.ast.Visitor;

import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

public class TypeAndRefChecker implements Visitor {

	Scope program_scope;
	Set<Node> initialized = new HashSet<Node>();
	DeclClass curr_class = null;
	DeclMethod curr_method = null;

	private int lastFoundMainLine = -1; // THERE CAN ONLY BE ONE!!!
	private boolean flgStatic = false;

	Map<String, String> key_inherits_from_value = new HashMap<String, String>();

	private boolean isGoodMainCandidate(DeclStaticMethod m) {
		if (!m.getName().equals("main"))
			return false;

		if (!m.getType().equals(new PrimitiveType(m.getLine(), DataType.VOID)))
			throw new LanguageException(m.getLine(),
					"Main method should have 'void' return type");

		PrimitiveType arg = new PrimitiveType(m.getLine(), DataType.STRING);
		arg.incrementDimension();

		if (m.getFormals().size() != 1
				|| !m.getFormals().get(0).getType().equals(arg))
			throw new LanguageException(m.getLine(),
					"Argument for main method should be 'string[] args'");

		return true;

	}

	private DeclClass checkClassName(String classname, Node curr_node) {
		Node res = program_scope.getFlat(classname);
		if (res == null)
			throw new LanguageException(curr_node.getLine(), "class "
					+ classname + " not defined for this scope");
		else
			return (DeclClass) res;
	}

	private Type checkTypeRef(HasType ht) {
		Type t = ht.getType();
		if (t instanceof ClassType) {
			// No need to check ref for primitive type
			ClassType ct = (ClassType) t;
			checkClassName(ct.getClassName(), (Node) ht);
		}

		return t;
	}

	private void multiapply(Collection<? extends Node> col) {
		for (Node node : col) {
			node.accept(this);
		}
	}

	@Override
	public Object visit(Program program) {
		program_scope = program.getScope();

		for (DeclClass c : program.getClasses()) {
			if (c.hasSuperClass()) {
				key_inherits_from_value.put(c.getName(), c.getSuperClassName());
			}
		}

		for (DeclClass c : program.getClasses()) {
			c.accept(this);
		}

		if (lastFoundMainLine < 0) {
			throw new LanguageException(program.getLine(),
					"no main defined for program");
		}

		return null;
	}

	@Override
	public Object visit(DeclClass c) {
		curr_class = c;
		multiapply(c.getFields());
		multiapply(c.getMethods());
		curr_class = null;
		return null;
	}

	@Override
	public Object visit(DeclField field) {
		return checkTypeRef(field);
	}

	@Override
	public Object visit(DeclVirtualMethod method) {
		checkTypeRef(method);
		curr_method = method;
		multiapply(method.getFormals());
		multiapply(method.getStatements());
		curr_method = null;
		return method.getType();
	}

	@Override
	public Object visit(DeclStaticMethod method) {
		checkTypeRef(method);
		curr_method = method;
		flgStatic = true;
		multiapply(method.getFormals());
		multiapply(method.getStatements());
		flgStatic = false;
		curr_method = null;

		if (isGoodMainCandidate(method)) {
			if (lastFoundMainLine >= 0) {
				throw new LanguageException(method.getLine(),
				/*
				 * "main for this program already defined in line " +
				 * lastFoundMainLine);
				 */
				"Found more than one main in the file");
			} else
				lastFoundMainLine = method.getLine();
		}

		return method.getType();
	}

	@Override
	public Object visit(DeclLibraryMethod method) {
		return method.getType();
	}

	@Override
	public Object visit(Parameter formal) {
		return checkTypeRef(formal);
	}

	@Override
	public Object visit(PrimitiveType type) {
		return type;
	}

	@Override
	public Object visit(ClassType type) {
		return type;
	}

	private boolean flgEnforceInitRef = true;

	@Override
	public Object visit(StmtAssignment a) {
		// System.out.println(a.getVariable() + ", " + a.getAssignment());
		flgEnforceInitRef = false;
		Object vt = a.getVariable().accept(this);
		flgEnforceInitRef = true;
		Object at = a.getAssignment().accept(this);

		if (at.equals(vt)) {
			Ref av = a.getVariable();
			if (av instanceof RefVariable) {
				Node v = a.getVariable().getScope()
						.getRecursive(((RefVariable) av).getName());
				initialized.add(v);
			}

			return vt;
		} else
			throw new LanguageException(a.getLine(),
					"Invalid assignment of type "
							+ ((Type) at).getDisplayName() + " to "
							+ ((vt instanceof RefField) ? "field" : "variable")
							+ " of type " + ((Type) vt).getDisplayName());
	}

	@Override
	public Object visit(StmtCall cs) {
		return cs.getCall().accept(this);
	}

	@Override
	public Object visit(StmtReturn rs) {
		if (rs.hasValue()) {
			Type rt = (Type) rs.getValue().accept(this);
			if (!rt.equals(curr_method.getType())) {
				throw new LanguageException(rs.getLine(),
						"return statement is not of type "
								+ curr_method.getType().getDisplayName());
			}
		} else {
			if (!(curr_method.getType() instanceof PrimitiveType)
					|| ((PrimitiveType) curr_method.getType()).getType() != DataType.VOID) {
				throw new LanguageException(rs.getLine(),
						"must return a value for a non-void method");
			}
		}

		return null;
	}

	@Override
	public Object visit(StmtIf is) {
		Type t = (Type) is.getCondition().accept(this);

		if (!t.equals(new PrimitiveType(is.getLine(), DataType.BOOLEAN))) {
			throw new LanguageException(is.getLine(),
					"Non boolean condition for if statement");
		}

		is.getOperation().accept(this);
		if (is.hasElse())
			is.getElseOperation().accept(this);
		return null;
	}

	@Override
	public Object visit(StmtWhile ws) {
		Type t = (Type) ws.getCondition().accept(this);

		if (!t.equals(new PrimitiveType(ws.getLine(), DataType.BOOLEAN))) {
			throw new LanguageException(ws.getLine(),
					"Non boolean condition for while statement");
		}

		ws.getOperation().accept(this);
		return null;
	}

	@Override
	public Object visit(StmtBreak breakStatement) {
		return null;
	}

	@Override
	public Object visit(StmtContinue s) {
		return null;
	}

	@Override
	public Object visit(StmtBlock b) {
		multiapply(b.getStatements());
		return null;
	}

	@Override
	public Object visit(LocalVariable lv) {
		if (lv.isInitialized()) {
			final Type init_type = (Type) lv.getInitialValue().accept(this);
			boolean flgGoodNull = false;
			if (!isGoodType(lv.getType(), init_type)) {
				PrimitiveType mock_null = new PrimitiveType(0, DataType.VOID);
				if (init_type.equals(mock_null)) {
					if (lv.getType() instanceof ClassType)
						flgGoodNull = true;
					else {
						PrimitiveType pt = (PrimitiveType) lv.getType();
						if (pt.getArrayDimension() > 0
								|| pt.getType().equals(DataType.STRING)) {
							flgGoodNull = true;
						}
					}
				}
				if (!flgGoodNull)
					throw new LanguageException(lv.getLine(),
							"initial value must be of same target type");
			}
			if (!flgGoodNull)
				initialized.add(lv);
		}

		return null;
	}

	private boolean isGoodType(Type lv, Type init_type) {
		if (lv.equals(init_type))
			return true;

		// check inheritance
		if (lv instanceof ClassType && init_type instanceof ClassType) {
			boolean res = isLeftInheirtsFromRight(
					((ClassType) init_type).getClassName(),
					((ClassType) lv).getClassName());
			return res;
		} else
			return false;
	}

	private boolean isLeftInheirtsFromRight(String left, String right) {
		String super_class = key_inherits_from_value.get(left);
		if (super_class == null)
			return false;
		if (super_class.equals(right))
			return true;
		return isLeftInheirtsFromRight(super_class, right);
	}

	@Override
	public Object visit(RefVariable r) {
		Node n = r.getScope().getRecursive(r.getName());

		// space for avoiding auto-capitalization
		if (n == null) {
			if (flgStatic && r.getScope().getParent() != null
					&& r.getScope().getParent().owner instanceof DeclClass) {
				DeclClass owner = (DeclClass) r.getScope().getParent().owner;

				if (owner.getInstanceScope().getRecursive(r.getName()) != null) {
					throw new LanguageException(r.getLine(),
							"Use a of field inside a static method is not allowed");
				}
			}

			throw new LanguageException(r.getLine(), " " + r.getName()
					+ " not found in symbol table");
		}

		if (!(n instanceof LocalVariable || n instanceof Parameter || n instanceof DeclField)) {
			throw new LanguageException(r.getLine(), r.getName()
					+ " refer to non-variable");
		}

		if (flgEnforceInitRef && !(n instanceof Parameter)
				&& !(n instanceof DeclField)
				&& ((HasType) n).getType().getArrayDimension() == 0
				&& !initialized.contains(n)) {
			throw new LanguageException(r.getLine(),
					"unable to refer to non initialized variable");
		}

		return ((HasType) n).getType();
	}

	private boolean initialized_contains_byname(HasName n) {
		for (Node in : initialized) {
			if (n.getName().equals(((HasName) in).getName())) {
				return true;
			}
		}

		return false;
	}

	@Override
	public Object visit(RefField r) {
		Node n;
		if (r.getObject() != null) {
			Type t = (Type) r.getObject().accept(this);
			if (!(t instanceof ClassType) || t.getArrayDimension() != 0) {
				throw new LanguageException(r.getLine(),
						"object must be a class instance");
			}

			final String classname = ((ClassType) t).getClassName();
			DeclClass c = checkClassName(classname, r);
			r.getObject().setExtraTypeInformation(t); // XXX for later target
														// evaluation in
														// TacVisitor
			n = c.getInstanceScope().getRecursive(r.getField());
		} else
			n = r.getScope().getRecursive(r.getField());

		if (n == null)
			throw new LanguageException(r.getLine(), "field " + r.getField()
					+ " was not found in scope");

		if (!(n instanceof DeclField)) {
			throw new LanguageException(r.getLine(), r.getField()
					+ " does not refer to a field");
		}

		if (false && flgEnforceInitRef && !initialized.contains(n))
			throw new LanguageException(n.getLine(),
					"unable to refer to non initialized field");

		return ((HasType) n).getType();
	}

	@Override
	public Object visit(RefArrayElement l) {
		Type array = (Type) l.getArray().accept(this);
		Type index = (Type) l.getIndex().accept(this);

		if (checkSimpleInt(index)) {
			// good index

			return array.cloneWithLowerDimension();
		} else {
			throw new LanguageException(l.getLine(),
					"index must be a scalar int");
		}
	}

	private boolean checkSimpleInt(Type index) {
		return index instanceof PrimitiveType
				&& ((PrimitiveType) index).getType() == DataType.INT
				&& index.getArrayDimension() == 0;
	}

	private void validateArgumentsForMethod(DeclMethod m,
			List<Expression> args, Node call) {
		final List<Parameter> parameters = m.getFormals();

		if (parameters.size() != args.size())
			throw new LanguageException(call.getLine(),
					"Invalid number of arguments for "
							+ ((DeclClass) m.getScope().owner).getName() + "."
							+ m.getName());

		for (int i = 0; i < args.size(); ++i) {
			Expression arg = args.get(i);
			Parameter p = parameters.get(i);

			Type ta = (Type) arg.accept(this);
			Type tp = (Type) p.accept(this);

			if (!isGoodType(tp, ta)) {
				throw new LanguageException(call.getLine(), "Method "
						+ ((DeclClass) m.getScope().owner).getName() + "."
						+ m.getName()
						+ " is not applicable for the arguments given");
			}
		}
	}

	@Override
	public Object visit(StaticCall call) {
		final String classname = call.getClassName();
		final String methodname = call.getMethod();
		final List<Expression> args = call.getArguments();

		DeclClass c = checkClassName(classname, call);
		Node n = c.getStaticScope().getRecursive(methodname);
		if (n == null
				|| !(n instanceof DeclStaticMethod || n instanceof DeclLibraryMethod)) {
			throw new LanguageException(call.getLine(), "Method " + methodname
					+ " doesn't exist");
		}

		DeclMethod m = (DeclMethod) n;

		validateArgumentsForMethod(m, args, call);

		call.setExtraTypeInformation(m.getType());

		return m.getType();
	}

	@Override
	public Object visit(VirtualCall call) {
		final String methodname = call.getMethod();

		Node n;
		if (call.hasExplicitObject()) {
			Type t = (Type) call.getObject().accept(this);
			if (!(t instanceof ClassType) || t.getArrayDimension() != 0) {
				throw new LanguageException(call.getLine(),
						"object must be a class instance");
			}

			call.getObject().setExtraTypeInformation(t); // XXX for later
															// evaluation at
															// TacVisitor
															// handling of
															// VirtualCall
			final String classname = ((ClassType) t).getClassName();
			DeclClass c = checkClassName(classname, call);
			n = c.getInstanceScope().getRecursive(methodname);
		} else {
			n = call.getScope().getRecursive(methodname);

			if (flgStatic) {
				if (n != null && !(n instanceof DeclStaticMethod)) {
					throw new LanguageException(call.getLine(),
							"Calling a local virtual method from inside a static method is not allowed");
				}
			}
		}

		if (n == null
		/*
		 * || !(n instanceof DeclVirtualMethod && !flgStatic || flgStatic && n
		 * instanceof DeclStaticMethod)
		 */) {
			// XXX
			throw new LanguageException(call.getLine(), " " + methodname
					+ " not found in symbol table");
		}

		DeclMethod m = (DeclMethod) n;

		validateArgumentsForMethod(m, call.getArguments(), call);

		call.setExtraTypeInformation(m.getType());

		return m.getType();
	}

	@Override
	public Object visit(This te) {
		if (flgStatic)
			throw new LanguageException(te.getLine(),
					"Use of 'this' expression inside static method is not allowed");

		return new ClassType(te.getLine(), curr_class.getName());
	}

	@Override
	public Object visit(NewInstance nc) {
		final String classname = nc.getName();

		checkClassName(classname, nc);

		return new ClassType(nc.getLine(), classname);
	}

	@Override
	public Object visit(NewArray na) {
		if (!checkSimpleInt((Type) na.getSize().accept(this))) {
			throw new LanguageException(na.getLine(), "size is not an int");
		}

		na.getType().incrementDimension();

		return na.getType();
	}

	@Override
	public Object visit(Length length) {
		Type arr = (Type) length.getArray().accept(this);
		if (arr.getArrayDimension() == 0)
			throw new LanguageException(length.getLine(),
					"unable to check length of scalar");
		return new PrimitiveType(length.getLine(), DataType.INT);
	}

	@Override
	public Object visit(Literal l) {
		return new PrimitiveType(l.getLine(), l.getType());
	}

	@Override
	public Object visit(UnaryOp op) {
		Type t = (Type) op.getOperand().accept(this);
		boolean flgBadType = false;
		if (t.getArrayDimension() == 0 && t instanceof PrimitiveType) {
			PrimitiveType pt = (PrimitiveType) t;
			switch (op.getOperator()) {
			case LNEG:
				flgBadType = pt.getType() != DataType.BOOLEAN;
				break;
			case UMINUS:
				flgBadType = pt.getType() != DataType.INT;
				break;
			}
		} else
			flgBadType = true;

		if (flgBadType)
			throw new LanguageException(op.getLine(),
					errorMessageForUnaryOp(op));

		return t;
	}

	@Override
	public Object visit(BinaryOp op) {
		Type op1 = (Type) op.getFirstOperand().accept(this);
		op.getFirstOperand().setExtraTypeInformation(op1);
		Type op2 = (Type) op.getSecondOperand().accept(this);
		op.getSecondOperand().setExtraTypeInformation(op2);

		PrimitiveType mock_null = new PrimitiveType(0, DataType.VOID);

		if (!op1.equals(op2) && !op1.equals(mock_null)
				&& !op2.equals(mock_null)) {
			throw new LanguageException(op.getLine(),
					errorMessageForBinaryOp(op));
		}

		switch (op.getOperator()) {
		// Equality comparison operators
		case NEQUAL:
		case EQUAL:
			return new PrimitiveType(op.getLine(), DataType.BOOLEAN);
		default:
			break;
		}

		if (op1.getArrayDimension() != 0 || !(op1 instanceof PrimitiveType)) {
			throw new LanguageException(op.getLine(),
					errorMessageForBinaryOp(op));
		}

		PrimitiveType pt = (PrimitiveType) op1;

		boolean flgMayBeString = false;
		switch (op.getOperator()) {
		// Arithmetic operators
		case PLUS:
			if (pt.getType() == DataType.STRING)
				return pt;
			flgMayBeString = true;
		case MINUS:
		case MOD:
		case MULTIPLY:
		case DIVIDE:
			if (pt.getType() == DataType.INT)
				return pt;
			else {
				throw new LanguageException(op.getLine(),
						errorMessageForBinaryOp(op));
			}
		default:
			break;
		}

		switch (op.getOperator()) {
		// Conditional operators
		case LAND:
		case LOR:
			if (pt.getType() == DataType.BOOLEAN)
				return pt;
			else {
				throw new LanguageException(op.getLine(),
						errorMessageForBinaryOp(op));
			}
		default:
			break;
		}

		switch (op.getOperator()) {
		// Relational comparison operators
		case GT:
		case GTE:
		case LT:
		case LTE:
			if (pt.getType() == DataType.INT)
				return new PrimitiveType(op.getLine(), DataType.BOOLEAN);
			else {
				throw new LanguageException(op.getLine(),
						errorMessageForBinaryOp(op));
			}
		default:
			break;

		}

		throw new LanguageException(op.getLine(),
				"types inappropriate to any operator");
	}

	private static String errorMessageForBinaryOp(BinaryOp op) {
		boolean flgOnlyInts = false;
		boolean flgOnlySameType = false;
		boolean flgOnlyBoolean = false;
		boolean flgOnlyIntsOrStrings = false;

		switch (op.getOperator()) {
		case DIVIDE:
		case GT:
		case GTE:
		case LT:
		case LTE:
		case MINUS:
		case MOD:
		case MULTIPLY:
			flgOnlyInts = true;
			break;
		case EQUAL:
		case NEQUAL:
			flgOnlySameType = true;
			break;
		case PLUS:
			flgOnlyIntsOrStrings = true;
			break;
		case LAND:
		case LOR:
			flgOnlyBoolean = true;
			break;
		}

		String postfix = "";

		if (flgOnlyInts) {
			postfix = "on non-integer expression";
		} else if (flgOnlySameType) {
			postfix = "on expressions of different types";
		} else if (flgOnlyIntsOrStrings) {
			postfix = "on non-integer and non-string expression";
		} else if (flgOnlyBoolean) {
			postfix = "on non-boolean expression";
		}

		postfix = " " + postfix;

		return "Invalid logical binary op (" + op.getOperator().toString()
				+ ")" + postfix;
	}

	private static String errorMessageForUnaryOp(UnaryOp op) {
		String required_type = "";

		switch (op.getOperator()) {
		case LNEG:
			required_type = "boolean";
			break;
		case UMINUS:
			required_type = "integer";
			break;
		}

		return "Invalid logical binary op (" + op.getOperator().toString()
				+ ") on non-" + required_type + " expression";
	}
}
